// Copyright (c) 2024 Go-Frame-Lite
// Go-Frame-Lite is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package crypto

import (
	"crypto/rand"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"time"

	"github.com/tjfoc/gmsm/sm2"
	gm_x509 "github.com/tjfoc/gmsm/x509"
)

// GenerateSM2KeyPair 生成SM2密钥对
func GenerateSM2KeyPair() (*sm2.PrivateKey, *sm2.PublicKey, error) {
	privateKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, fmt.Errorf("生成SM2密钥对失败: %w", err)
	}

	return privateKey, &privateKey.PublicKey, nil
}

// GenerateSM2Certificate 生成SM2证书
func GenerateSM2Certificate(subject pkix.Name, validityDays int,
	privateKey *sm2.PrivateKey, parentCert *gm_x509.Certificate, parentKey *sm2.PrivateKey) ([]byte, error) {

	template := gm_x509.Certificate{
		SerialNumber:          big.NewInt(time.Now().Unix()),
		Subject:               subject,
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(0, 0, validityDays),
		KeyUsage:              gm_x509.KeyUsageKeyEncipherment | gm_x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []gm_x509.ExtKeyUsage{gm_x509.ExtKeyUsageServerAuth, gm_x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	var certBytes []byte
	var err error

	if parentCert == nil || parentKey == nil {
		// 自签名证书
		template.IsCA = true
		template.KeyUsage |= gm_x509.KeyUsageCertSign
		certBytes, err = gm_x509.CreateCertificate(&template, &template, &privateKey.PublicKey, privateKey)
	} else {
		// 由父证书签名
		certBytes, err = gm_x509.CreateCertificate(&template, parentCert, &privateKey.PublicKey, parentKey)
	}

	if err != nil {
		return nil, fmt.Errorf("生成SM2证书失败: %w", err)
	}

	return certBytes, nil
}

// GenerateSM2KeyPairAndCertificate 生成SM2密钥对和自签名证书
func GenerateSM2KeyPairAndCertificate(subject pkix.Name, validityDays int) ([]byte, []byte, error) {
	privateKey, _, err := GenerateSM2KeyPair()
	if err != nil {
		return nil, nil, err
	}

	certBytes, err := GenerateSM2Certificate(subject, validityDays, privateKey, nil, nil)
	if err != nil {
		return nil, nil, err
	}

	// 将私钥编码为PEM格式
	privateKeyBytes, err := marshalSM2PrivateKey(privateKey)
	if err != nil {
		return nil, nil, fmt.Errorf("编码SM2私钥失败: %w", err)
	}

	privateKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "SM2 PRIVATE KEY",
		Bytes: privateKeyBytes,
	})

	// 将证书编码为PEM格式
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	})

	return privateKeyPEM, certPEM, nil
}

// GenerateSM2CertificateChain 生成SM2证书链（根证书、中间证书、终端证书）
func GenerateSM2CertificateChain(rootSubject, intermediateSubject, endEntitySubject pkix.Name, validityDays int) (map[string][]byte, error) {
	result := make(map[string][]byte)

	// 生成根证书
	rootPrivateKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("生成根证书密钥失败: %w", err)
	}

	rootCertBytes, err := GenerateSM2Certificate(rootSubject, validityDays*3, rootPrivateKey, nil, nil)
	if err != nil {
		return nil, err
	}

	// 生成中间证书
	intermediatePrivateKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("生成中间证书密钥失败: %w", err)
	}

	rootCert, err := gm_x509.ParseCertificate(rootCertBytes)
	if err != nil {
		return nil, fmt.Errorf("解析根证书失败: %w", err)
	}

	intermediateCertBytes, err := GenerateSM2Certificate(intermediateSubject, validityDays*2, intermediatePrivateKey, rootCert, rootPrivateKey)
	if err != nil {
		return nil, err
	}

	// 生成终端证书
	endEntityPrivateKey, err := sm2.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("生成终端证书密钥失败: %w", err)
	}

	intermediateCert, err := gm_x509.ParseCertificate(intermediateCertBytes)
	if err != nil {
		return nil, fmt.Errorf("解析中间证书失败: %w", err)
	}

	endEntityCertBytes, err := GenerateSM2Certificate(endEntitySubject, validityDays, endEntityPrivateKey, intermediateCert, intermediatePrivateKey)
	if err != nil {
		return nil, err
	}

	// 编码所有密钥和证书
	result["root_private_key"] = encodePrivateKeyToPEM(rootPrivateKey)
	result["root_certificate"] = encodeCertificateToPEM(rootCertBytes)
	result["intermediate_private_key"] = encodePrivateKeyToPEM(intermediatePrivateKey)
	result["intermediate_certificate"] = encodeCertificateToPEM(intermediateCertBytes)
	result["end_entity_private_key"] = encodePrivateKeyToPEM(endEntityPrivateKey)
	result["end_entity_certificate"] = encodeCertificateToPEM(endEntityCertBytes)

	return result, nil
}

// encodePrivateKeyToPEM 将SM2私钥编码为PEM格式
func encodePrivateKeyToPEM(privateKey *sm2.PrivateKey) []byte {
	privateKeyBytes, err := marshalSM2PrivateKey(privateKey)
	if err != nil {
		return nil
	}

	return pem.EncodeToMemory(&pem.Block{
		Type:  "SM2 PRIVATE KEY",
		Bytes: privateKeyBytes,
	})
}

// encodeCertificateToPEM 将证书编码为PEM格式
func encodeCertificateToPEM(certBytes []byte) []byte {
	return pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	})
}

// ParseSM2PrivateKey 解析PEM格式的SM2私钥
func ParseSM2PrivateKey(privateKeyPEM []byte) (*sm2.PrivateKey, error) {
	block, _ := pem.Decode(privateKeyPEM)
	if block == nil {
		return nil, fmt.Errorf("解析PEM块失败")
	}

	// 解析私钥
	privateKey, err := parseSM2PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("解析SM2私钥失败: %w", err)
	}

	return privateKey, nil
}

// ParseSM2Certificate 解析PEM格式的SM2证书
func ParseSM2Certificate(certPEM []byte) (*gm_x509.Certificate, error) {
	block, _ := pem.Decode(certPEM)
	if block == nil {
		return nil, fmt.Errorf("解析PEM块失败")
	}

	cert, err := gm_x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("解析SM2证书失败: %w", err)
	}

	return cert, nil
}

// marshalSM2PrivateKey 将SM2私钥编码为ASN.1格式
func marshalSM2PrivateKey(privateKey *sm2.PrivateKey) ([]byte, error) {
	// 简单的私钥编码实现
	privateKeyBytes := privateKey.D.Bytes()

	// 如果私钥长度小于32字节，用0填充
	if len(privateKeyBytes) < 32 {
		padded := make([]byte, 32)
		copy(padded[32-len(privateKeyBytes):], privateKeyBytes)
		privateKeyBytes = padded
	}

	return privateKeyBytes, nil
}

// parseSM2PrivateKey 从ASN.1格式解析SM2私钥
func parseSM2PrivateKey(data []byte) (*sm2.PrivateKey, error) {
	// 简单的私钥解析实现
	if len(data) < 32 {
		return nil, fmt.Errorf("私钥数据长度不足")
	}

	// 提取私钥值
	d := new(big.Int).SetBytes(data)

	// 创建私钥对象
	privateKey := new(sm2.PrivateKey)
	privateKey.D = d
	privateKey.PublicKey.Curve = sm2.P256Sm2()
	privateKey.PublicKey.X, privateKey.PublicKey.Y = privateKey.PublicKey.Curve.ScalarBaseMult(d.Bytes())

	return privateKey, nil
}
